//===-- RegAllocSpillAll.cpp - spill all register allocator ---------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This file defines the RegisterAllocator function pass, which provides a minimal
// implementation of the spill all register allocator.
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "spiller"
#include "llvm/Function.h"
#include "llvm/ADT/BitVector.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/CodeGen/LiveVariables.h"
#include "llvm/CodeGen/MachineFrameInfo.h"
#include "llvm/CodeGen/MachineFunction.h"
#include "llvm/CodeGen/MachineFunctionPass.h"
#include "llvm/CodeGen/MachineInstr.h"
#include "llvm/CodeGen/MachineLoopInfo.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
#include "llvm/CodeGen/Passes.h"
#include "llvm/CodeGen/RegAllocRegistry.h"
#include "llvm/Target/TargetInstrDesc.h"
#include "llvm/Target/TargetInstrInfo.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/Target/TargetOptions.h"
#include "llvm/Target/TargetRegisterInfo.h"
#include "llvm/Support/GraphWriter.h"
#include "llvm/Support/raw_ostream.h"

#include <map>
#include <set>
#include <vector>
#include <stack>

using namespace llvm;

static RegisterRegAlloc spillAllRegAlloc("spillall", "spill all register allocator", createSpillAllRegisterAllocator);

static RegisterRegAlloc fullRegAlloc("full", "full register allocator", createFullRegisterAllocator);

STATISTIC(LoadSpillCount, "Number of loads spilled");
STATISTIC(StoreSpillCount, "Number of stores spilled");

/* Live Interval Assignment State */
typedef enum LIAssignment {
  UNREMOVED,            // A live interval that has not yet been removed from the interference graph (this is the initial state of every node)
  REMOVED_PHYS_REG,     // A live interval that has been removed from the interference graph and not marked as a spill candidate
  REMOVED_SPILL_CAND,   // A live interval that has been removed from the interference graph and marked as a spill candidate
  PHYS_REG,             // The live interval that has been returned to the interference graph and assigned to a physical registered
  SPILLED               // The live interval that has been returned to the interference graph and assigned to a memory location on the stack
} LIAssignment;

/* Live Interval */
class LiveInterval {
    /* Private member variables */
    LIAssignment assignment; // Indicates the current state of a live interval
    unsigned virtualReg; // The virtual register corresponding to this live interval
    std::set<LiveInterval*> neighbors; // Set of LiveIntervals that this LiveInterval interferes with (i.e. neighbors in an interference graph)
    unsigned unremovedNeighborCount; // The number of neighbors that have not yet been removed from the interference graph
    union {
      int physReg; // The physical register the live range is assigned to (assignment = PHYS_REG)
      int spillOffset; // The offset from the function pointer of the memory location the live range is assigned to (assignment = SPILL)
    };

    /* Private methods */
    void notifyNeighborRemoved(); // Notify live interval that one of its neighbors was removed from interference graph
    void notifyNeighborReturned(); // Notify live interval that one of its neighbors was returned to interference graph

  public:
    /* Constructor */
    LiveInterval();
    LiveInterval(unsigned virtualReg);

    /* State Readers */
    bool isRemoved();
    bool isSpillCandid();
    bool isAssignedToPhysReg();
    bool isSpilled();

    /* Attribute Readers */
    int getVirtualReg() const;
    int getPhysReg();
    int getSpillOffset();
    unsigned getUnremovedNeighborCount();

    /* Methods used when constructing the interferance graph */
    void addNeighbor(LiveInterval* neighbor); // Add interfering LiveInterval to this interval's neighbors set

    /* Methods used during graph coloring - remove phase */
    void remove(); // Remove this live interval from the interference graph
    void removeAndMarkSpill(); // Remove this live interval from the interference graph and mark it as a spill candidate

    /* Methods used during graph coloring - return phase */
    void assignToPhysReg(int physReg); // Return the live interval to the graph and assign it to a physical register
    void spillToOffset(int offset); // Return the live interval to the graph and assign it to a stack memory location

    const std::set<LiveInterval *> &getNeighbors() { return neighbors; }
};

LiveInterval::LiveInterval() {
  this->assignment = UNREMOVED;
  this->unremovedNeighborCount = 0;
  this->virtualReg = -1;
}
LiveInterval::LiveInterval(unsigned virtualReg) {
  this->assignment = UNREMOVED;
  this->unremovedNeighborCount = 0;
  this->virtualReg = virtualReg;
}

bool LiveInterval::isRemoved() { return (assignment == REMOVED_PHYS_REG) || (assignment == REMOVED_SPILL_CAND); }
bool LiveInterval::isSpillCandid() { return assignment == REMOVED_SPILL_CAND; }
bool LiveInterval::isAssignedToPhysReg() { return assignment == PHYS_REG; }
bool LiveInterval::isSpilled() { return assignment == SPILLED; }

int LiveInterval::getVirtualReg() const { return virtualReg; }
int LiveInterval::getPhysReg() { return physReg; }
int LiveInterval::getSpillOffset() { return spillOffset; }
unsigned LiveInterval::getUnremovedNeighborCount() { return unremovedNeighborCount; }

void LiveInterval::addNeighbor(LiveInterval* neighbor) {
  if (neighbors.insert(neighbor).second)
    unremovedNeighborCount++;
}

void LiveInterval::remove() {
  this->assignment = REMOVED_PHYS_REG;
  for (std::set<LiveInterval*>::iterator it = neighbors.begin(); it != neighbors.end(); it++) {
    (*it)->notifyNeighborRemoved(); // Notify them that it has been removed
  }
}

void LiveInterval::removeAndMarkSpill() {
  this->assignment = REMOVED_SPILL_CAND;
  for (std::set<LiveInterval*>::iterator it = neighbors.begin(); it != neighbors.end(); it++) {
    (*it)->notifyNeighborRemoved(); // Notify them that it has been removed
  }
}

void LiveInterval::notifyNeighborRemoved() { unremovedNeighborCount--; }

void LiveInterval::assignToPhysReg(int physReg) {
  this->assignment = PHYS_REG;
  this->physReg = physReg;
  for (std::set<LiveInterval*>::iterator it = neighbors.begin(); it != neighbors.end(); it++) {
    (*it)->notifyNeighborReturned(); // Notify them that it has been returned
  }
}

void LiveInterval::spillToOffset(int offset) {
  this->assignment = SPILLED;
  this->spillOffset = offset;
  for (std::set<LiveInterval*>::iterator it = neighbors.begin(); it != neighbors.end(); it++) {
    (*it)->notifyNeighborReturned(); // Notify them that it has been returned
  }
}

void LiveInterval::notifyNeighborReturned() { unremovedNeighborCount++; }

namespace {

class RegisterAllocator : public MachineFunctionPass
{
  MachineFunction *MF; // Context
  std::vector<LiveInterval*> allLIs; // Set of LiveIntervals
  std::map<unsigned,LiveInterval*> regLIMap; // A map from virtual registers to live intervals
  std::map< MachineBasicBlock* , std::set<LiveInterval*>* > bbToLiveOutMap; // A map from a basic blocks to its LiveOut set

  bool forceSpills; // SpillAll or Full? Gets us to maximize code reuse
public:
  RegisterAllocator(bool spillAll);

  /// Return the pass name.
  virtual const char* getPassName() const {
    return forceSpills ? "Spill All Register Allocator"
      : "Full Register Allocator";
  }

  /// RegisterAllocator analysis usage.
  virtual void getAnalysisUsage(AnalysisUsage &AU) const;

  /// Perform register allocation.
  virtual bool runOnMachineFunction(MachineFunction &mf);

  // Member functions to build live intervals and the interference graph
  // Note that live intervals include both physical and virtual registers, to
  // better handle cases like function calls and return value interferences.
  void buildLiveIntervals();
  void buildLiveOutSets();
  void buildLiveIntervalInterferences();
  void resetLiveIntervals();

  // Actual Chaitin-Briggs graph coloring
  float estimateSpillCost(LiveInterval* LI);
  bool colorGraph();
  LiveInterval* getLiveIntervalForOperand(MachineOperand &OP);

  // Insert spill code (into virtual registers)
  void insertSpills();
  // Make virtual registers physical
  void rewriteRegisters();
  static char ID;
};

char RegisterAllocator::ID = 0;

} // end anonymous namespace

RegisterAllocator::RegisterAllocator(bool spillAll): MachineFunctionPass(ID), forceSpills(spillAll) {
  initializeUnreachableBlockElimPass(*PassRegistry::getPassRegistry());
  initializeLiveVariablesPass(*PassRegistry::getPassRegistry());
  initializeStrongPHIEliminationPass(*PassRegistry::getPassRegistry());
}

void RegisterAllocator::getAnalysisUsage(AnalysisUsage &AU) const {
  AU.setPreservesCFG();
  AU.addRequiredID(StrongPHIEliminationID);
  AU.addRequired<LiveVariables>();
  AU.addRequired<MachineLoopInfo>();
  MachineFunctionPass::getAnalysisUsage(AU);
}

bool RegisterAllocator::runOnMachineFunction(MachineFunction &mf) {

  MF = &mf;

  // Because of our architecture, the only difference between the spill-all
  // allocator and the full architecture is that spill-all gets to first mark
  // everybody as spilled.
  if (forceSpills) {
    resetLiveIntervals();
    buildLiveIntervals();
    // Inform all live intervals that they need to be spilled.
    MachineFrameInfo *MFI = MF->getFrameInfo();
    for (unsigned liIndex = 0; liIndex < allLIs.size(); liIndex++ ) {
      allLIs[liIndex]->spillToOffset(MFI->CreateSpillStackObject(4, 4));
    }
    // Spill 'em all!
    insertSpills();
  }
  // While we can't color the graph, spill vregs that need to be spilled.
  do {
    resetLiveIntervals();
    buildLiveIntervals();
    buildLiveOutSets();
    buildLiveIntervalInterferences();
    if (!colorGraph())
      break;
    insertSpills();
  } while (true);

  // Our graph is now fully colored. Rewrite registers and call it done!
  rewriteRegisters();

  return true;
}

void RegisterAllocator::rewriteRegisters() {
  // Actually rewrite all virtual registers with physical registers.
  const TargetRegisterInfo *TRI = MF->getTarget().getRegisterInfo();
  MachineRegisterInfo &MRI = MF->getRegInfo();
  for (MachineFunction::iterator BB = MF->begin(), E = MF->end(); BB != E;
      ++BB) {
    for (MachineBasicBlock::iterator MI = BB->begin(), MIE = BB->end();
        MI != MIE; ++MI) {
      for (unsigned op = 0; op < MI->getNumOperands(); op++) {
        MachineOperand &MO = MI->getOperand(op);
        if (!MO.isReg())
          continue;
        LiveInterval *LI = getLiveIntervalForOperand(MO);
        if (LI) {
          // By the time we reach this function, all live intervals should be
          // assigned to a physical register. If it hasn't, we're screwed
          if (!LI->isAssignedToPhysReg())
            return;
          MO.substPhysReg(LI->getPhysReg(), *TRI);
          MRI.setPhysRegUsed(LI->getPhysReg());
        }
      }
    }
  }
}

float RegisterAllocator::estimateSpillCost(LiveInterval* LI) {
  const TargetInstrInfo *TII = MF->getTarget().getInstrInfo();
  // Compute cost
  int cost = 0;
  int virtualReg = LI->getVirtualReg();
  for(MachineRegisterInfo::reg_iterator OP = MF->getRegInfo().reg_begin(virtualReg); OP != MachineRegisterInfo::reg_end(); OP++) { // For every reference of the live interval
    int loopDepth = getAnalysis<MachineLoopInfo>().getLoopDepth(OP->getParent()); // Get the instruction's loop depth
    int loopDepthFactor = pow(10, loopDepth); // Compute loop depth factor: loopDepthFactor = 10^loopDepth
    TargetInstrDesc desc = TII->get(OP->getOpcode());
    // Since we don't load/store saving of spills, we don't decrease those
    // values, so we just always count 1 for a ref.
    cost += loopDepthFactor;
  }

  // Compute degree
  int degree = LI->getUnremovedNeighborCount();

  // Spill cost = cost / degree
  return static_cast<float>(cost)/static_cast<float>(degree);

}

void RegisterAllocator::insertSpills() {
  // This will insert all of the spill code for any register that needs to be
  // spilled.
  const TargetInstrInfo *TII = MF->getTarget().getInstrInfo();
  const TargetRegisterInfo *TRI = MF->getTarget().getRegisterInfo();
  MachineRegisterInfo &MRI = MF->getRegInfo();
  for (MachineFunction::iterator BB = MF->begin(), E = MF->end(); BB != E;
      ++BB) {
    int skip = 0;
    for (MachineBasicBlock::iterator MI = BB->begin(), MIE = BB->end();
        MI != MIE; ++MI) {
      // Skip over our generated spill code
      if (skip > 0) {
        skip--;
        continue;
      }
      // Get the registers to spill for the current instruction.
      for (unsigned op = 0; op < MI->getNumOperands(); op++) {
        // Grab the live interval for the register, and make sure it is valid.
        MachineOperand &MO = MI->getOperand(op);
        LiveInterval *spillee = getLiveIntervalForOperand(MO);
        if (!spillee || !spillee->isSpilled())
          continue;
        if (!MO.isReg() || !TargetRegisterInfo::isVirtualRegister(MO.getReg()))
          continue;

        // Okay, now use a new virtual register for spilling (we'll color again,
        // so this won't be an issue).
        const TargetRegisterClass *regClass =
          MRI.getRegClass(MO.getReg());
        int newReg = MRI.createVirtualRegister(regClass);
        if (MO.isDef()) {
          // Definitions = store after the instruction
          // I really wish that this were a random-access iterator...
          MachineBasicBlock::iterator target = MI;
          ++skip;
          for (int i = 0; i < skip; i++)
            ++target;
          TII->storeRegToStackSlot(*BB, target, newReg, true,
            spillee->getSpillOffset(), regClass, TRI);
          StoreSpillCount++;
        } else {
          // Uses = store before the instruction
          TII->loadRegFromStackSlot(*BB, MI, newReg, spillee->getSpillOffset(),
            regClass, TRI);
          LoadSpillCount++;
        }
        // And replace as appropriate
        MO.setReg(newReg);
      }
    }
  }
}

LiveInterval* RegisterAllocator::getLiveIntervalForOperand(MachineOperand &OP) {
  if (OP.isReg() && regLIMap.count(OP.getReg()) ) {
    return regLIMap.find(OP.getReg())->second;
  } else {
    return NULL;
  }
}

void RegisterAllocator::buildLiveOutSets() {
  // Seed the live outs for the function for physical registers.
  {
    MachineFunction::iterator MBB = MF->begin();
    MachineRegisterInfo &MRI = MF->getRegInfo();
    std::set<LiveInterval*> *liveout = new std::set<LiveInterval*>();
    for (MachineRegisterInfo::liveout_iterator I = MRI.liveout_begin(),
        E = MRI.liveout_end(); I != E; ++I) {
      liveout->insert(regLIMap[*I]);
    }
    for (MachineFunction::iterator E = MF->end(); MBB != E; ++MBB) {
      if (MBB->succ_empty()) {
        bbToLiveOutMap[MBB] = liveout;
      }
    }
  }

  for (MachineFunction::iterator BB = MF->begin(); BB != MF->end(); BB++) { // For every basic blocks in the function
    std::set<LiveInterval*> *liveOut = new std::set<LiveInterval*>();
    for (std::vector<LiveInterval*>::iterator LI = allLIs.begin(); LI != allLIs.end(); ++LI) { // For every live interval
      unsigned reg = (*LI)->getVirtualReg(); // Get the virtual register assigned to the live interval
      if (!TargetRegisterInfo::isVirtualRegister(reg))
        continue;
      if (getAnalysis<LiveVariables>().isLiveOut (reg, *BB)) { // If the virtual register is live at the output of the block
        liveOut->insert(*LI); // Add the corresponding live interval to the live out set
      }
    }
    bbToLiveOutMap.insert ( std::pair< MachineBasicBlock* , std::set<LiveInterval*>* >(BB,liveOut) ); // Add the generated live out set to the map
  }
}

void RegisterAllocator::buildLiveIntervalInterferences() {

  for (MachineFunction::iterator BB = MF->begin(); BB != MF->end(); BB++) { // For every basic blocks in the function
    std::set<LiveInterval*> *liveNow = new std::set<LiveInterval*>();
    *liveNow = *(bbToLiveOutMap[BB]); // Initialize the liveNow set to liveOut
    for (MachineBasicBlock::iterator MI = BB->end(); MI != BB->begin(); ) { // Walk backwards through the instructions of the block
      MI--;
      // All defs interfere with liveNow and are removed from liveNow
      for (MachineInstr::mop_iterator OP = MI->operands_begin(), E = MI->operands_end(); OP != E; ++OP) { // For every operand in the instruction
        if (!OP->isReg()) continue;
        if (OP->isDef()) { // If the operand is a def
          LiveInterval *LI = getLiveIntervalForOperand(*OP); // Get its corresponding live interval
          if (LI) {
            // So, there appears to be a few cases where we have some implicit
            // gunk where explicit def-use chains don't work well. We'll still
            // mark it as interfering for robustness, but it isn't necessarily
            // be in the liveNow set.
            std::set<LiveInterval*>::iterator it = liveNow->find(LI);
            if (it != liveNow->end())
              liveNow->erase(it);
            for (std::set<LiveInterval*>::iterator LINow = liveNow->begin(); LINow != liveNow->end(); LINow++) { // For every live interval in live now
              // Set them to interfere with each other
              LI->addNeighbor(*LINow);
              (*LINow)->addNeighbor(LI);
            }
          }
        }
      }
      // All uses are added from liveNow
      for (MachineInstr::mop_iterator OP = MI->operands_begin(), E = MI->operands_end(); OP != E; ++OP) { // For every operand in the instruction
        if (!OP->isReg()) continue;
        if (OP->isUse()) { // If this is a use...
          LiveInterval* LI = getLiveIntervalForOperand(*OP); // Get its corresponding live interval
          if (LI)
            liveNow->insert(LI); // Add it to live now
        }
      }
    }
  }
}

void RegisterAllocator::buildLiveIntervals() {
  // Map every virtual register and physical register to a separate live interval
  for (MachineFunction::iterator BB = MF->begin(); BB != MF->end(); BB++) { // For every basic blocks in the function
    for (MachineBasicBlock::iterator MI = BB->begin(); MI != BB->end(); MI++) { // For every instruction in the basic block
      for (MachineInstr::mop_iterator OP = MI->operands_begin(), E = MI->operands_end(); OP != E; ++OP) { // For every operand in the instruction
        if (OP->isReg()) { // If the operand is a virtual or physical register
          if ( ! regLIMap.count(OP->getReg()) ) { // If the virtual register has not yet been mapped to a live interval
            LiveInterval* LI = new LiveInterval(OP->getReg()); // Create a new live interval from the virtual register
            regLIMap.insert( std::pair<unsigned,LiveInterval*>(OP->getReg(),LI) ); // Map the virtual register to the live interval

            // If it's a physical register, be sure to note that fact.
            if (TargetRegisterInfo::isPhysicalRegister(OP->getReg()))
              LI->assignToPhysReg(OP->getReg());
          }
        }
      }
    }
  }

  // Add the final set of live intervals to the allLIs
  for (std::map<unsigned,LiveInterval*>::iterator it = regLIMap.begin(); it != regLIMap.end(); it++) {
    allLIs.push_back(it->second);
  }
}

void RegisterAllocator::resetLiveIntervals() {
  // Clear out everything we built for live intervals and interferences, so that
  // we can start another round of the graph coloring algorithm
  allLIs.clear();
  regLIMap.clear();
  bbToLiveOutMap.clear();
}

bool RegisterAllocator::colorGraph() {
  bool insertedSpills = false;
  std::stack<LiveInterval*> liStack; // Intermediate stack to place element in while running the coloring algorithm
  const TargetRegisterInfo *TRI = MF->getTarget().getRegisterInfo();
  MachineRegisterInfo &MRI = MF->getRegInfo();

  // Phase 0: Collect all physical registers from the map, and remove from
  // further consideration. They're already allocated, so no sense trying to
  // figure out how to color them here.
  std::vector<LiveInterval*> physRegs;
  for (std::vector<LiveInterval*>::iterator it = allLIs.begin();
      it != allLIs.end(); ++it) {
    if ((*it)->isAssignedToPhysReg()) {
      std::vector<LiveInterval*>::iterator save = it - 1;
      physRegs.push_back(*it);
      allLIs.erase(it);
      it = save;
    }
  }

  // Phase 1: removing elements from the graph to the stack
  while (allLIs.size() > 0) { // Iterate until all graph elements have been removed
    bool noElementsToRemove = true; // Flag to track if any elements were removed on this iteration; if none have been, we will need to spill a node
    for (unsigned liIndex = 0; liIndex < allLIs.size(); ) { // Traverse list of live intervals in the graph
      LiveInterval *LI = allLIs[liIndex];

      // How many neighbors can it have?
      const TargetRegisterClass *RC = MRI.getRegClass(LI->getVirtualReg());
      BitVector availableRegs = TRI->getAllocatableSet(*MF, RC);
      if (LI->getUnremovedNeighborCount() < availableRegs.count()) { // If the number of remaining neighbors is less than that of physical registers
        LI->remove(); // Mark the node as removed
        liStack.push(LI); // Push it on the stack
        allLIs.erase(allLIs.begin() + liIndex); // Remove it from the graph
        noElementsToRemove = false; // Flag that an element has been removed on this iteration
      } else {
        // Proceed to the next one
        liIndex++;
      }
    }
    if (noElementsToRemove) { // If no elements have been removed on this iteration
      // Find the live interval with the lowest spill cost
      unsigned minIndex = 0;
      float minSpillCost = estimateSpillCost(allLIs[minIndex]);
      for (unsigned liIndex = 1; liIndex < allLIs.size(); liIndex++) { // Traverse list of live intervals in the graph
        float liSpillCost = estimateSpillCost(allLIs[liIndex]); // Get the current interval's spill cost
        if (liSpillCost < minSpillCost) { // If the current interval has lower spill cost than the minimum
          // It becomes the minimum
          minIndex = liIndex;
          minSpillCost = liSpillCost;
        }
      }

      // Use it
      allLIs[minIndex]->removeAndMarkSpill(); // Mark a node as a spill candidate
      liStack.push(allLIs[minIndex]); // Push it on the stack
      allLIs.erase(allLIs.begin() + minIndex); // Remove it from the graph
    }
  }

  // Phase 2: returning elements from the stack to the graph
  MachineFrameInfo *MFI = MF->getFrameInfo();

  while (liStack.size() > 0) {
    LiveInterval* LI = liStack.top(); // Get the topmost element on the stack
    liStack.pop(); // Pop it from the stack
    allLIs.push_back(LI); // Return it to the graph

    // Which registers are available? Get the full set, and exclude any that
    // have already been allocated by neighbors.
    const TargetRegisterClass *RC = MRI.getRegClass(LI->getVirtualReg());
    BitVector availableRegs = TRI->getAllocatableSet(*MF, RC);
    for (std::set<LiveInterval *>::const_iterator it =
        LI->getNeighbors().begin(); it != LI->getNeighbors().end(); ++it) {
      if ((*it)->isAssignedToPhysReg())
        availableRegs.reset((*it)->getPhysReg());
    }
    int availablePhysReg = availableRegs.find_first();
    if (availablePhysReg >= 0) { // If an available physical register has been found
      LI->assignToPhysReg(availablePhysReg); // Mark as returned to graph and assign to available physical register
    } else { // If no physical register is available
      LI->spillToOffset(MFI->CreateSpillStackObject(4, 4)); // Spill it
      insertedSpills = true;
    }
  }
  return insertedSpills;
}

FunctionPass* llvm::createSpillAllRegisterAllocator() {
  return new RegisterAllocator(true);
}

FunctionPass* llvm::createFullRegisterAllocator() {
  return new RegisterAllocator(false);
}
